{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimizing Stiffness Constants\n",
    "\n",
    "Suppose we have $n$ blocks, each of width $w$, at positions $x\\in\\mathbb{R}^n$.\n",
    "The blocks are positioned between two walls at $0$ and $l$.\n",
    "The leftmost block is connected to the left wall via a spring with stiffness coefficient $k_1$,\n",
    "the rightmost block is connected to the right wall via a spring with stiffness coefficient $k_{n+1}$,\n",
    "and the $i$th block is connected to the $(i+1)$th block via a spring with stiffness coefficient $k_{i+1}$.\n",
    "The equilibrium position of all the blocks can be found by solving the optimization problem\n",
    "\\begin{equation}\n",
    "\\begin{array}{ll}\n",
    "\\mbox{minimize} & \\frac{1}{2}k_1x_1^2 + \\frac{1}{2}k_2(x_2-x_1)^2 + \\ldots + \\frac{1}{2}k_{n+1}(l-x_n)^2\\\\[.2cm]\n",
    "\\mbox{subject to} & x_1 \\geq w/2, \\quad x_n \\leq l - w/2, \\\\\n",
    "& x_i - x_{i-1} \\geq w, \\quad i=2,\\ldots,n-1,\n",
    "\\end{array}\n",
    "\\label{eq:prob}\n",
    "\\end{equation}\n",
    "with variable $x$ and solution denoted $x^\\star$.\n",
    "The objective is the potential energy of the system, and the constraints express\n",
    "the fact that the blocks have a width, and cannot penetrate each other or the walls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: this notebook requires ffmpeg, which can be installed on ubuntu with \"sudo apt install ffmpeg\"\n",
    "# and on mac with \"brew install ffmpeg\"\n",
    "import cvxpy as cp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from cvxpylayers.torch import CvxpyLayer\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib import animation, rc\n",
    "rc('animation', html='html5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can easily set up this problem as a CVXPY problem, with parameter $k$ and variable $x$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "l = 1\n",
    "w = .05\n",
    "\n",
    "k = cp.Parameter(n + 1, nonneg=True)\n",
    "x = cp.Variable(n + 2)\n",
    "objective = cp.sum(cp.multiply(k, .5 * cp.square(cp.diff(x))))\n",
    "constraints = [x[0] == 0, x[-1] == l] + [x[i] - x[i-1] >= w for i in range(1, n+2)]\n",
    "prob = cp.Problem(cp.Minimize(objective), constraints)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can easily convert this problem to a `torch` `CvxpyLayer` in one line.\n",
    "The layer maps stiffness coefficients to block positions, or $x^\\star(k)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob_tch = CvxpyLayer(prob, [k], [x])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our goal is to tune the stiffness coefficients such that the block positions are close to target positions,\n",
    "or solve the problem\n",
    "\\begin{equation}\n",
    "\\begin{array}{ll}\n",
    "\\mbox{minimize} & \\|x^\\star(k) - x^\\mathrm{targ}\\|_2^2,\n",
    "\\end{array}\n",
    "\\end{equation}\n",
    "with variable $k\\in\\mathbb{R}^{n+1}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "k = torch.ones(n+1, dtype=torch.double, requires_grad=True)\n",
    "x_targ, _ = torch.sort(torch.rand(n, dtype=torch.double))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can (approximately) solve this problem via gradient descent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, = prob_tch(k)\n",
    "x_np = x.detach().numpy()\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_xlim((0, 1))\n",
    "ax.set_ylim((0, w*4))\n",
    "rects = []\n",
    "for i in range(n):\n",
    "    plt.axvline(x_targ[i].item(), c='k')\n",
    "    rect = patches.Rectangle((x_np[i + 1] - w/2, 0), w, w, linewidth=1, edgecolor='k', facecolor='none')\n",
    "    ax.add_patch(rect)\n",
    "    rects.append(rect)\n",
    "\n",
    "def animate(i):\n",
    "    k.grad = None\n",
    "    x, = prob_tch(k)\n",
    "    x_np = x.detach().numpy()\n",
    "    loss = (x[1:-1] - x_targ).pow(2).sum()\n",
    "    loss.backward()\n",
    "    k.data = k.data - .1 * k.grad.data\n",
    "    k.data[k.data < 0] = 0.0\n",
    "    for i in range(n):\n",
    "        rects[i].xy = (x_np[i + 1] - w/2, 0)\n",
    "    return rects\n",
    "\n",
    "anim = animation.FuncAnimation(fig, animate,\n",
    "                               frames=300, interval=50, blit=True)\n",
    "anim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
